import math
import torch
import copy
import torch.nn.functional as F
from torch import nn
import numpy as np

from .inference import make_fcos_postprocessor
from .loss import make_align_loss_evaluator

from taa_core.layers import Scale
from taa_core.layers import DeformConv, DFConv2d
from .utils import ConvBlock

INF = 1e8


class FCOSHead(torch.nn.Module):
    def __init__(self, cfg, 
        in_channels, 
        regress_ranges=((-1, 64), (64, 128), (128, 256), (256, 512),
                (512, INF))
    ):
        """
        Arguments:
            in_channels (int): number of channels of the input feature
        """
        super(FCOSHead, self).__init__()
        num_classes = cfg.MODEL.FCOS.NUM_CLASSES - 1
        self.fpn_strides = cfg.MODEL.FCOS.FPN_STRIDES
        self.norm_reg_targets = cfg.MODEL.FCOS.NORM_REG_TARGETS
        self.centerness_on_reg = cfg.MODEL.FCOS.CENTERNESS_ON_REG
        self.use_dcn_in_tower = cfg.MODEL.FCOS.USE_DCN_IN_TOWER
        self.bbox_norm_type = cfg.MODEL.FCOS.BBOX_NORM_TYPE
        self.reg_denoms = [
            regress_range[-1] for regress_range in regress_ranges
        ]

        self.gradient_mul = 0.1
        self.num_dconv_points = 9
        dcn_kernel = int(math.sqrt(self.num_dconv_points))
        dcn_pad = int((dcn_kernel - 1) / 2)
        dcn_base = np.arange(-dcn_pad,
                            dcn_pad + 1).astype(np.float64)
        dcn_base_y = np.repeat(dcn_base, dcn_kernel)
        dcn_base_x = np.tile(dcn_base, dcn_kernel)
        dcn_base_offset = np.stack([dcn_base_y, dcn_base_x], axis=1).reshape((-1))
        self.dcn_base_offset = torch.tensor(dcn_base_offset).reshape(1, -1, 1, 1)

        cls_tower = []
        bbox_tower = []
        com_tower = []
        assert cfg.MODEL.FCOS.NUM_CONVS == 3
        for i in range(cfg.MODEL.FCOS.NUM_CONVS):
            if self.use_dcn_in_tower and \
                    i == cfg.MODEL.FCOS.NUM_CONVS - 1:
                conv_func = DFConv2d
            else:
                conv_func = nn.Conv2d

            cls_tower.append(
                conv_func(
                    in_channels,
                    in_channels,
                    kernel_size=3,
                    stride=1,
                    padding=1,
                    bias=True
                )
            )
            cls_tower.append(nn.GroupNorm(32, in_channels))
            cls_tower.append(nn.ReLU())
            bbox_tower.append(
                conv_func(
                    in_channels,
                    in_channels,
                    kernel_size=3,
                    stride=1,
                    padding=1,
                    bias=True
                )
            )
            bbox_tower.append(nn.GroupNorm(32, in_channels))
            bbox_tower.append(nn.ReLU())
            com_tower.append(ConvBlock(in_channels, in_channels, kernel_size=3, padding=1, 
                norm_func=nn.GroupNorm(32, in_channels)))

        com_tower.append(ConvBlock(in_channels, in_channels, kernel_size=3, padding=1, 
            norm_func=nn.GroupNorm(32, in_channels)))
        self.add_module('cls_tower', nn.Sequential(*cls_tower))
        self.add_module('bbox_tower', nn.Sequential(*bbox_tower))
        self.add_module('com_tower', nn.Sequential(*com_tower))

        # self.bbox_pred_conv = ConvBlock(
        #     in_channels, in_channels, kernel_size=3, padding=1, 
        #     norm_func=nn.GroupNorm(32, in_channels)
        # )
        self.bbox_pred = nn.Conv2d(
            in_channels, 4, kernel_size=3, stride=1,
            padding=1
        )

        self.vfnet_reg_refine_dconv = DeformConv(
            in_channels,
            in_channels,
            dcn_kernel,
            1,
            padding=dcn_pad)
        self.vfnet_reg_refine = nn.Conv2d(in_channels, 4, 3, padding=1)

        self.vfnet_cls_dconv = DeformConv(
            in_channels,
            in_channels,
            dcn_kernel,
            1,
            padding=dcn_pad)
        self.cls_logits = nn.Conv2d(
            in_channels, num_classes, kernel_size=3, stride=1,
            padding=1
        )

        self.centerness = nn.Conv2d(
            in_channels, 1, kernel_size=3, stride=1, padding=1)

        # initialization
        for modules in [
            self.cls_tower, self.bbox_tower, 
            self.com_tower, self.cls_logits, self.bbox_pred,
            self.vfnet_reg_refine_dconv, self.vfnet_reg_refine, 
            self.vfnet_cls_dconv, self.centerness
        ]:
            for layer in modules.modules():
                if isinstance(layer, nn.Conv2d):
                    torch.nn.init.normal_(layer.weight, mean=0, std=0.01)
                    torch.nn.init.constant_(layer.bias, 0)
                if isinstance(layer, nn.GroupNorm):
                    torch.nn.init.constant_(layer.weight, 1)
                    torch.nn.init.constant_(layer.bias, 0)

        # initialize the bias for focal loss
        prior_prob = cfg.MODEL.FCOS.PRIOR_PROB
        bias_value = -math.log((1 - prior_prob) / prior_prob)
        torch.nn.init.constant_(self.cls_logits.bias, bias_value)
        torch.nn.init.constant_(self.centerness.bias, bias_value)

        self.scales = nn.ModuleList([Scale(init_value=1.0) for _ in range(5)])
        self.scale_refines = nn.ModuleList([Scale(init_value=1.0) for _ in range(5)])

    def forward(self, x):
        logits = []
        bbox_reg = []
        centerness = []
        bbox_pred_refines = []

        for l, feature in enumerate(x):
            cls_tower = self.cls_tower(feature)
            box_tower = self.bbox_tower(feature)
            com_tower = self.com_tower(feature)

            # bbox_pred = self.scales[l](self.bbox_pred(com_tower))
            if self.bbox_norm_type == 'reg_denom':
                bbox_pred = self.scales[l](
                    self.bbox_pred(com_tower)).float().exp() * self.reg_denoms[l]
            elif self.bbox_norm_type == 'stride':
                bbox_pred = self.scales[l](
                    self.bbox_pred(com_tower)).float().exp() * self.fpn_strides[l]
            else:
                bbox_pred = self.scales[l](self.bbox_pred(com_tower)).exp()
            bbox_reg.append(bbox_pred)

            centerness.append(self.centerness(com_tower))

            dcn_offset = self.star_dcn_offset(bbox_pred, self.gradient_mul,
                                            self.fpn_strides[l]).to(box_tower.dtype)

            box_tower = F.relu(self.vfnet_reg_refine_dconv(box_tower, dcn_offset))
            bbox_pred_refine = self.scale_refines[l](
                self.vfnet_reg_refine(box_tower)).float().exp()
            bbox_pred_refine = bbox_pred_refine * bbox_pred.detach()
            bbox_pred_refines.append(bbox_pred_refine)

            cls_tower = F.relu(self.vfnet_cls_dconv(cls_tower, dcn_offset))
            logits.append(self.cls_logits(cls_tower))           

        return logits, bbox_reg, bbox_pred_refines, centerness

    def star_dcn_offset(self, bbox_pred, gradient_mul, stride):
        """Compute the star deformable conv offsets.

        Args:
            bbox_pred (Tensor): Predicted bbox distance offsets (l, r, t, b).
            gradient_mul (float): Gradient multiplier.
            stride (int): The corresponding stride for feature maps,
                used to project the bbox onto the feature map.

        Returns:
            dcn_offsets (Tensor): The offsets for deformable convolution.
        """
        dcn_base_offset = self.dcn_base_offset.type_as(bbox_pred)
        bbox_pred_grad_mul = (1 - gradient_mul) * bbox_pred.detach() + \
            gradient_mul * bbox_pred
        # map to the feature map scale
        bbox_pred_grad_mul = bbox_pred_grad_mul / stride
        N, C, H, W = bbox_pred.size()

        x1 = bbox_pred_grad_mul[:, 0, :, :]
        y1 = bbox_pred_grad_mul[:, 1, :, :]
        x2 = bbox_pred_grad_mul[:, 2, :, :]
        y2 = bbox_pred_grad_mul[:, 3, :, :]
        bbox_pred_grad_mul_offset = bbox_pred.new_zeros(
            N, 2 * self.num_dconv_points, H, W)
        bbox_pred_grad_mul_offset[:, 0, :, :] = -1.0 * y1  # -y1
        bbox_pred_grad_mul_offset[:, 1, :, :] = -1.0 * x1  # -x1
        bbox_pred_grad_mul_offset[:, 2, :, :] = -1.0 * y1  # -y1
        bbox_pred_grad_mul_offset[:, 4, :, :] = -1.0 * y1  # -y1
        bbox_pred_grad_mul_offset[:, 5, :, :] = x2  # x2
        bbox_pred_grad_mul_offset[:, 7, :, :] = -1.0 * x1  # -x1
        bbox_pred_grad_mul_offset[:, 11, :, :] = x2  # x2
        bbox_pred_grad_mul_offset[:, 12, :, :] = y2  # y2
        bbox_pred_grad_mul_offset[:, 13, :, :] = -1.0 * x1  # -x1
        bbox_pred_grad_mul_offset[:, 14, :, :] = y2  # y2
        bbox_pred_grad_mul_offset[:, 16, :, :] = y2  # y2
        bbox_pred_grad_mul_offset[:, 17, :, :] = x2  # x2
        dcn_offset = bbox_pred_grad_mul_offset - dcn_base_offset

        return dcn_offset


class FCOSModule(torch.nn.Module):
    """
    Module for FCOS computation. Takes feature maps from the backbone and
    FCOS outputs and losses. Only Test on FPN now.
    """

    def __init__(self, cfg, in_channels):
        super(FCOSModule, self).__init__()

        head = FCOSHead(cfg, in_channels)

        box_selector_test = make_fcos_postprocessor(cfg)

        loss_evaluator = make_align_loss_evaluator(cfg)
        self.head = head
        self.box_selector_test = box_selector_test
        self.loss_evaluator = loss_evaluator
        self.fpn_strides = cfg.MODEL.FCOS.FPN_STRIDES

    def forward(self, images, features, targets=None):
        """
        Arguments:
            images (ImageList): images for which we want to compute the predictions
            features (list[Tensor]): features computed from the images that are
                used for computing the predictions. Each tensor in the list
                correspond to different feature levels
            targets (list[BoxList): ground-truth boxes present in the image (optional)

        Returns:
            boxes (list[BoxList]): the predicted boxes from the RPN, one BoxList per
                image.
            losses (dict[Tensor]): the losses for the model during training. During
                testing, it is an empty dict.
        """
        (box_cls, 
         box_regression, 
         box_regression_refine, 
         centerness) = self.head(features)

        locations = self.compute_locations(features)
 
        if self.training:
            return self._forward_train(
                locations=locations, 
                box_cls=box_cls, 
                box_regression=box_regression, 
                box_regression_refine=box_regression_refine,
                centerness=centerness,
                targets=targets
            )
        else:
            return self._forward_test(
                locations, 
                box_cls=box_cls,
                box_regression=box_regression,
                box_regression_refine=box_regression_refine,
                centerness=centerness,
                image_sizes=images.image_sizes
            )

    def _forward_train(self, locations, box_cls, box_regression, 
                        box_regression_refine, centerness, targets
    ):
        losses = self.loss_evaluator(
            locations, 
            box_cls=box_cls, 
            box_regression=box_regression,
            box_regression_refine=box_regression_refine,
            centerness=centerness,
            targets=targets
        )

        return None, losses

    def _forward_test(self, locations, box_cls, box_regression, 
        box_regression_refine, centerness, image_sizes
    ):
        boxes = self.box_selector_test(
            locations, box_cls, box_regression, 
            box_regression_refine, centerness, image_sizes
        )
        return boxes, {}

    def compute_locations(self, features):
        locations = []
        for level, feature in enumerate(features):
            h, w = feature.size()[-2:]
            locations_per_level = self.compute_locations_per_level(
                h, w, self.fpn_strides[level],
                feature.device
            )
            locations.append(locations_per_level)
        return locations

    def compute_locations_per_level(self, h, w, stride, device):
        shifts_x = torch.arange(
            0, w * stride, step=stride,
            dtype=torch.float32, device=device
        )
        shifts_y = torch.arange(
            0, h * stride, step=stride,
            dtype=torch.float32, device=device
        )
        shift_y, shift_x = torch.meshgrid(shifts_y, shifts_x)
        shift_x = shift_x.reshape(-1)
        shift_y = shift_y.reshape(-1)
        locations = torch.stack((shift_x, shift_y), dim=1) + stride // 2
        return locations

def build_align(cfg, in_channels):
    return FCOSModule(cfg, in_channels)
